package com.sunday.common.mq.kafka.study.spring.e20_Brave_Trace;

import brave.Tracing;
import brave.kafka.clients.KafkaTracing;
import brave.messaging.MessagingTracing;
import com.sunday.common.mq.kafka.brave.core.TracingKafkaConsumerFactory;
import com.sunday.common.mq.kafka.brave.core.TracingKafkaProducerFactory;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.kafka.*;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
import org.springframework.kafka.config.ContainerCustomizer;
import org.springframework.kafka.config.KafkaListenerContainerFactory;
import org.springframework.kafka.core.ConsumerFactory;
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;

import java.util.Map;

@Slf4j
@Configuration
@RequiredArgsConstructor
public class KafkaBraveConfiguration {

    private final KafkaProperties properties;

    @Bean
    public MessagingTracing messagingTracing(Tracing tracing) {
        return MessagingTracing.create(tracing);
    }

    @Bean
    public KafkaTracing springRocketTracing(MessagingTracing messagingTracing) {
        return KafkaTracing.create(messagingTracing);
    }

    /**
     * Producer 埋点
     */
    @Bean
    public TracingKafkaProducerFactory<?, ?> kafkaProducerFactory(KafkaConnectionDetails connectionDetails, ObjectProvider<DefaultKafkaProducerFactoryCustomizer> customizers, KafkaTracing kafkaTracing) {
        Map<String, Object> properties = this.properties.buildProducerProperties();
        properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, connectionDetails.getProducerBootstrapServers());
        TracingKafkaProducerFactory<?, ?> factory = new TracingKafkaProducerFactory<>(properties, kafkaTracing);
        String transactionIdPrefix = this.properties.getProducer().getTransactionIdPrefix();
        if (transactionIdPrefix != null) {
            factory.setTransactionIdPrefix(transactionIdPrefix);
        }
        customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
        return factory;
    }


    /**
     * consumer 埋点
     */
    @Bean
    public TracingKafkaConsumerFactory<?, ?> kafkaConsumerFactory(KafkaConnectionDetails connectionDetails, ObjectProvider<DefaultKafkaConsumerFactoryCustomizer> customizers, KafkaTracing kafkaTracing) {
        Map<String, Object> properties = this.properties.buildConsumerProperties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, connectionDetails.getConsumerBootstrapServers());
        TracingKafkaConsumerFactory<?, ?> factory = new TracingKafkaConsumerFactory<>(properties, kafkaTracing);
        customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
        return factory;
    }

    /**
     * https://docs.spring.io/spring-kafka/docs/current/reference/html/#batch-listeners
     */
    @Bean
    public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Object, Object>> kafkaBatchListenerContainerFactory(
            ConcurrentKafkaListenerContainerFactoryConfigurer configurer,
            ObjectProvider<ConsumerFactory<Object, Object>> kafkaConsumerFactory,
            ObjectProvider<ContainerCustomizer<Object, Object, ConcurrentMessageListenerContainer<Object, Object>>> kafkaContainerCustomizer,
            KafkaProperties properties
    ) {
        ConcurrentKafkaListenerContainerFactory<Object, Object> factory = new ConcurrentKafkaListenerContainerFactory<>();
        configurer.configure(factory, kafkaConsumerFactory
                .getIfAvailable(() -> new DefaultKafkaConsumerFactory<>(properties.buildConsumerProperties())));
        kafkaContainerCustomizer.ifAvailable(factory::setContainerCustomizer);
        factory.setBatchListener(true);
        return factory;
    }

}
